
import torch
import torch.nn as nn
import pybullet as p
import time
import pybullet_data
import random
import onnxruntime as ort
import onnx

if torch.cuda.is_available():
    torch_device = torch.device('cuda')
else:
    torch_device = torch.device('cpu')

startPos = [0, 0, 0.55]
motorsIdList = [0, 1, 2, 4, 5, 6, 8, 9, 10, 12, 13, 14]
testTorques = [0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0]
initialJointsPos = [0, 0.6, -1.2, 0, 0.6, -1.2, 0, 0.6, -1.2, 0, 0.6, -1.2]
pybulletTimeStep = 0.001
pos_gains=[20,20,20,20,20,20,20,20,20,20,20,20]
vel_gains=[0.5,0.5,0.5,0.5,0.5,0.5,0.5,0.5,0.5,0.5,0.5,0.5]
max_torques=[55,55,55,55,55,55,55,55,55,55,55,55]
count = 0
num = 0
zeros = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

# * Set pybullet environment
p.connect(p.GUI)
p.setGravity(0, 0, -9.81)
p.setTimeStep(pybulletTimeStep)
p.setAdditionalSearchPath(pybullet_data.getDataPath())  # optionally
planeId = p.loadURDF("plane.urdf")
quadrupedId = p.loadURDF("../urdf/lion/lion.urdf", startPos, useFixedBase=False)
# p.setRealTimeSimulation(1)
for i in range(p.getNumJoints(quadrupedId)):
    print(p.getJointInfo(quadrupedId, i))
print(p.getNumJoints(quadrupedId))


# * Set motors to position mode
for j in range(12):
    p.resetJointState(quadrupedId, motorsIdList[j], initialJointsPos[j])
for j in range(12):
    p.setJointMotorControl2(quadrupedId, motorsIdList[j], p.POSITION_CONTROL,force=0)
# * Set joint velocity limits
for j in range(16):
    p.changeDynamics(quadrupedId,j,maxJointVelocity=30,lateralFriction=100.0)


# print(p.getJointInfo(quadrupedId, 0))

# * Load actor network
path='../models/model_3000.pt'
class ActorCritic(nn.Module):
    def __init__(self):
        super(ActorCritic, self).__init__()
        num_obs = 48
        num_actions = 12
        # hidden_dims = [512,256,128]
        hidden_dims = [256,128,64]
        #Policy
        actor_layers = []
        actor_layers.append(nn.Linear(num_obs,hidden_dims[0]))
        actor_layers.append(nn.ELU())
        for l in range(len(hidden_dims)):
            if l==len(hidden_dims)-1:
                actor_layers.append(nn.Linear(hidden_dims[l],num_actions))
            else:
                actor_layers.append(nn.Linear(hidden_dims[l],hidden_dims[l+1]))
                actor_layers.append(nn.ELU())
        self.actor=nn.Sequential(*actor_layers)
        # Value function
        critic_layers = []
        critic_layers.append(nn.Linear(num_obs, hidden_dims[0]))
        critic_layers.append(nn.ELU())
        for l in range(len(hidden_dims)):
            if l == len(hidden_dims) - 1:
                critic_layers.append(nn.Linear(hidden_dims[l], 1))
            else:
                critic_layers.append(nn.Linear(hidden_dims[l], hidden_dims[l + 1]))
                critic_layers.append(nn.ELU())
        self.critic = nn.Sequential(*critic_layers)
        self.std = nn.Parameter(1.0 * torch.ones(num_actions))
        self.distribution = None
actor_critic=ActorCritic()
loaded_dict=torch.load(path, map_location=torch_device)
actor_critic.load_state_dict(loaded_dict['model_state_dict'])
actor_critic.eval()
# print(actor_critic)

# * Export model to onnx file
_input = torch.randn(48)
_output = actor_critic.actor(_input)
torch.onnx.export(actor_critic.actor,
                  _input,
                  'model.onnx',
                  opset_version=10,
                  do_constant_folding=True,
                  verbose=True,
                  input_names=['observations'],
                  output_names=['actions'],
                  )
# * Validate onnx model
ort_session = ort.InferenceSession('model.onnx')
loaded_model = onnx.load('model.onnx')
onnx.checker.check_model(loaded_model)
print(ort_session.get_inputs()[0].shape,ort_session.get_inputs()[0].name,ort_session.get_inputs()[0].type)


# * Update observations funcitons
class Obs:
    def __init__(self):
        self.joints_pos=torch.tensor(zeros)
        self.joints_vel=torch.tensor(zeros)
        self.base_pos=torch.tensor([0,0,0])
        self.base_vel=torch.tensor([0,0,0])
        self.base_omega=torch.tensor([0,0,0])
        self.base_quat=torch.tensor([0,0,0])
        self.actions=torch.tensor(zeros)

obs=Obs()
def q_conjugate(q):
    x, y, z, w = q
    return (-x, -y, -z, w)
def q_mult(q1, q2):
    x1, y1, z1, w1 = q1
    x2, y2, z2, w2 = q2
    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
    y = w1 * y2 + y1 * w2 + z1 * x2 - x1 * z2
    z = w1 * z2 + z1 * w2 + x1 * y2 - y1 * x2
    return x, y, z, w
# def quat_rotate_inverse(q1, v1):
#     q2 = v1 + (0.0,)
#     output = q_mult(q_mult(q1, q2), q_conjugate(q1))[:3]
#     return output
def quat_rotate_inverse(q1, v1):
    q2 = v1 + (0.0,)
    output = q_mult(q_mult(q_conjugate(q1), q2), q1)[:3]
    return output


def updateMotorsInfo():
    joint_states = p.getJointStates(quadrupedId, motorsIdList)
    pos_mid=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    vel_mid=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    for i in range(12):
        pos_mid[i] = joint_states[i][0]
        vel_mid[i] = joint_states[i][1]
    pos=pos_mid[3:6]+pos_mid[0:3]+pos_mid[9:12]+pos_mid[6:9]
    vel=vel_mid[3:6]+vel_mid[0:3]+vel_mid[9:12]+vel_mid[6:9]
    obs.joints_pos = torch.tensor(pos)
    obs.joints_vel = torch.tensor(vel)
    # print(obs.joints_vel)
    # print(torch.max(obs.joints_vel))

def updateIMUInfo():
    pos_ori = p.getBasePositionAndOrientation(quadrupedId)
    vel_omega = p.getBaseVelocity(quadrupedId)
    obs.base_pos = pos_ori[0]
    obs.base_vel = vel_omega[0]
    obs.base_omega = vel_omega[1]
    obs.base_quat = pos_ori[1]

def updateObs():
    updateIMUInfo()
    updateMotorsInfo()
    vel=quat_rotate_inverse(obs.base_quat, obs.base_vel)
    omega=quat_rotate_inverse(obs.base_quat,obs.base_omega)
    grav=quat_rotate_inverse(obs.base_quat,(0,0,-1)) # * These values should be right
    obs_buf=torch.cat((torch.tensor(vel)*1.0,
                       torch.tensor(omega)*1.,
                       torch.tensor(grav),
                       torch.tensor([.5,.0,0.]),
                       obs.joints_pos-torch.tensor(initialJointsPos),
                       obs.joints_vel*0.05,
                       obs.actions
                       ))
    return obs_buf

# p.resetBasePositionAndOrientation(quadrupedId,posObj=[0.,0.,0.5],ornObj=[  0.7071068, 0, 0, 0.7071068])

num_obs=(torch.zeros(48))
old_time=time.time()
decimation=20
for i in range(10000000):
    # p.resetBaseVelocity(quadrupedId, linearVelocity=[0, 0, 0], angularVelocity=[0, 0, 1])
    start=time.time()
    actions=actor_critic.actor(num_obs)
    stop=time.time()
    print((stop-start)*1000,"ms")
    num_obs = (torch.zeros(48))
    obs.actions = actions

    # print(actions.tolist())
    for _ in range(decimation):
        scaled_action=0.25*actions
        torques=20*(scaled_action+torch.tensor(initialJointsPos)-obs.joints_pos)-0.5*(obs.joints_vel)
        # print(obs.joints_vel)
        torch.clip(torques,-torch.tensor(max_torques),torch.tensor(max_torques))
        # command=torch.flip(torques,dims=[0])
        command_mid=torques.tolist()
        # print(command_mid[3:6]+command_mid[0:3]+command_mid[9:12]+command_mid[6:9])
        command=command_mid[3:6]+command_mid[0:3]+command_mid[9:12]+command_mid[6:9]
        # torques=30*(torch.tensor(initialJointsPos)-obs.joints_pos)-0.5*(obs.joints_vel)
        p.setJointMotorControlArray(quadrupedId, motorsIdList, p.TORQUE_CONTROL, forces=command)
        # p.setJointMotorControlArray(quadrupedId,motorsIdList,p.PD_CONTROL,targetPositions=(0.25*actions+torch.tensor(initialJointsPos)-obs.joints_pos),forces=zeros,targetVelocities=zeros,positionGains=pos_gains,velocityGains=vel_gains)
        p.stepSimulation()
        time.sleep(pybulletTimeStep)
        num_obs+=updateObs()
    # num_obs=updateObs()
    num_obs=num_obs*(1/decimation)
    # print(num_obs.tolist())

    # print("Period is: ",time.time()-old_time)
    old_time=time.time()




p.disconnect()
